#include <stdio.h>
#include <string.h>
#include "util.h"
#include "table.h"
#include "symbol.h"
#include "temp.h"
#include "graph.h"
#include "liveness.h"
#include "color.h"
#define MAX_NODE_NUM 1024

typedef struct G_nodeDList_ *G_nodeDList;
typedef struct Live_moveDList_ *Live_moveDList;

typedef enum {PRECOL, INITIAL, SIMPWORK, FREEWORK, SPILWORK, SPILNODE, COALNODE, COLNODE, STACK} nodeKind;
typedef enum {COAL, CONS, FROZ, WORK, ACT} moveKind;

static char* nodeStrs[10] = {"preColored", "initial", "simplifyWorkList", "freezeWorkList", "spillWorkList", "spilledNodes", "coalescedNodes", "coloredNodes", "selectStack"};
static char* moveStrs[10] = {"coalescedMoves", "constrainedMoves", "frozenMoves", "workListMoves", "activeMoves"};
static Temp_map debug = NULL; // only used for debug

typedef struct TAB_table_ *Move_table; // mapping from 'move' to 'moveDList'
 
struct G_nodeDList_ {
	G_node node;
	nodeKind kind;
	G_nodeDList prev;
	G_nodeDList next;
};

struct Live_moveDList_ {
	Live_move move;
	moveKind kind;
	Live_moveDList prev;
	Live_moveDList next;
};

static Live_moveDList Live_MoveDList(Live_move move, moveKind kind, Live_moveDList prev, Live_moveDList next) {
	Live_moveDList moveDList = checked_malloc(sizeof(*moveDList));
	moveDList->move = move;
	moveDList->kind = kind;
	moveDList->prev = prev;
	moveDList->next = next;
	return moveDList;
}

static G_nodeDList G_NodeDList(G_node node, nodeKind kind, G_nodeDList prev, G_nodeDList next) {
	G_nodeDList nodeDList = checked_malloc(sizeof(*nodeDList));
	nodeDList->node = node;
	nodeDList->kind = kind;
	nodeDList->prev = prev;
	nodeDList->next = next;
	return nodeDList;
}

static Move_table Move_empty() {
	return TAB_empty();
}

static void Move_enter(Move_table tab, Live_move key, Live_moveDList value) {
	TAB_enter(tab, key, value);
}

static Live_moveDList Move_look(Move_table tab, Live_move key) {
	return TAB_look(tab, key);
}	

static void insertNodeToSet(G_nodeDList list, G_nodeDList p) {
	// insert 'p' into 'list' and modify the kind of 'p' with that of 'list'
	G_nodeDList q = list->next;
	p->next = q;
	q->prev = p;
	p->prev = list;
	list->next = p;
	p->kind = list->kind;
	printf("insert %s to %s\n", Temp_look(debug, Live_gtemp(p->node)), nodeStrs[p->kind]);
} 

static void removeNodeFromSet(G_nodeDList p) {
	// remove 'p' from the list where 'p' live in
	G_nodeDList prev = p->prev;
	G_nodeDList next = p->next;
	prev->next = next;
	next->prev = prev;
	p->prev = p->next = NULL; 
	printf("remove %s from %s\n", Temp_look(debug, Live_gtemp(p->node)), nodeStrs[p->kind]);
}

static bool nodeSetEmpty(G_nodeDList list) {
	return !(list->next->next);
}

static void insertMoveToSet(Live_moveDList list, Live_moveDList p) {
	// insert 'p' into 'list' and modify the kind of 'p' with that of 'list'
	Live_moveDList q = list->next;
	p->next = q;
	q->prev = p;
	p->prev = list;
	list->next = p;
	p->kind = list->kind;
	printf("insert (%s, %s) to %s\n", Temp_look(debug, Live_gtemp(p->move->src)), Temp_look(debug, Live_gtemp(p->move->dst)), moveStrs[p->kind]);
}

static void removeMoveFromSet(Live_moveDList p) {
	// remove 'p' from the list where 'p' live in
	assert(p);
	Live_moveDList prev = p->prev;
	Live_moveDList next = p->next;
	prev->next = next;
	next->prev = prev;
	p->prev = p->next = NULL; 
	printf("remove (%s, %s) to %s\n", Temp_look(debug, Live_gtemp(p->move->src)), Temp_look(debug, Live_gtemp(p->move->dst)), moveStrs[p->kind]);
}

static bool moveSetEmpty(Live_moveDList list) {
	return !(list->next->next);
}

static G_nodeDList preColored = NULL;
static G_nodeDList initial = NULL;
static G_nodeDList simplifyWorkList = NULL;
static G_nodeDList freezeWorkList = NULL;
static G_nodeDList spillWorkList = NULL;
static G_nodeDList spilledNodes = NULL;
static G_nodeDList coalescedNodes = NULL;
static G_nodeDList coloredNodes = NULL;
static G_nodeDList selectStack = NULL; 

static Live_moveDList coalescedMoves = NULL;
static Live_moveDList constrainedMoves = NULL;
static Live_moveDList frozenMoves = NULL;
static Live_moveDList workListMoves = NULL;
static Live_moveDList activeMoves = NULL;

static int adjSet[MAX_NODE_NUM][MAX_NODE_NUM];
static G_nodeList adjList[MAX_NODE_NUM];
static int degree[MAX_NODE_NUM];
static Live_moveList moveList[MAX_NODE_NUM];
static int alias[MAX_NODE_NUM];
static Temp_temp color[MAX_NODE_NUM];
static G_node nodes[MAX_NODE_NUM]; // mapping from int to node

static Move_table moveHash = NULL; // mapping from 'move' to 'moveDList'
static G_table nodeHash = NULL; // mapping from 'node' to 'nodeDList'

static int K = 0; // the number of machine register

static void insertNodeToList(int nodeIndex, G_node node) {
	G_nodeList p = G_NodeList(node, adjList[nodeIndex]);
	adjList[nodeIndex] = p;
}

static void insertMoveToList(int nodeIndex, Live_move move) {
	Live_moveList p = Live_MoveList(move, moveList[nodeIndex]);
	moveList[nodeIndex] = p;
}

static void addEdge(G_node u, G_node v) {
	int uIndex = G_nodeIndex(u);
	int vIndex = G_nodeIndex(v);
	if(uIndex == vIndex || adjSet[uIndex][vIndex])
		return;
	adjSet[uIndex][vIndex] = 1;
	adjSet[vIndex][uIndex] = 1;
	G_nodeDList uDList = G_look(nodeHash, u);
	G_nodeDList vDList = G_look(nodeHash, v);
	assert(uDList && vDList);
	if(uDList->kind != PRECOL) {
		insertNodeToList(uIndex, v);
		++degree[uIndex];
	}
	if(vDList->kind != PRECOL) {
		insertNodeToList(vIndex, u);
		++degree[vIndex];
	}
	printf("%s interfere %s\n", Temp_look(debug, Live_gtemp(u)), Temp_look(debug, Live_gtemp(v)));
}

static G_nodeList adjacent(G_node node) {
	int nodeIndex = G_nodeIndex(node);
	G_nodeList dummy = G_NodeList(NULL, NULL);
	for(G_nodeList p = adjList[nodeIndex]; p; p = p->tail) {
		G_node neighbor = p->head;
		G_nodeDList q = G_look(nodeHash, neighbor);
		// adjList[node] / (selectStack U coalescedNodes)
		if(q->kind == STACK || q->kind == COALNODE)
			continue;
		dummy->tail = G_NodeList(neighbor, dummy->tail);
	}
	return dummy->tail;
}

static G_nodeList mergeNodeList(G_nodeList list1, G_nodeList list2) {
	// precondition: both elements in 'list1' and 'list2' are unique 
	G_table tab = G_empty();
	G_nodeList dummy = G_NodeList(NULL, NULL);
	for(G_nodeList p = list1; p; p = p->tail) {
		G_node node = p->head;
		G_enter(tab, node, (void*)1);		
		dummy->tail = G_NodeList(node, dummy->tail);
	}
	for(G_nodeList p = list2; p; p = p->tail) {
		G_node node = p->head;
		if(!G_look(tab, node))	
			dummy->tail = G_NodeList(node, dummy->tail);
	}
	return dummy->tail;
}

static Live_moveList nodeMoves(G_node node) {
	int nodeIndex = G_nodeIndex(node);
	Live_moveList dummy = Live_MoveList(NULL, NULL);
	for(Live_moveList p = moveList[nodeIndex]; p; p = p->tail) {
		Live_move move = p->head;
		Live_moveDList q = Move_look(moveHash, move);
		// moveList[node] ^ (activeMoves U workListMoves)
		if(q->kind != ACT && q->kind != WORK)	
			continue;
		dummy->tail = Live_MoveList(move, dummy->tail);			
	}
	return dummy->tail;
}

static bool moveRelated(G_node node) {
	return nodeMoves(node) != NULL;
}

static void makeWorkList() {
	G_nodeDList p = initial->next, q;
	while(p->next) {
		q = p->next;
		removeNodeFromSet(p);
		G_node node = p->node;
		int nodeIndex = G_nodeIndex(node);
		if(degree[nodeIndex] >= K) 
			insertNodeToSet(spillWorkList, p);
		else if(moveRelated(node))
			insertNodeToSet(freezeWorkList, p);
		else
			insertNodeToSet(simplifyWorkList, p);
		p = q;
	}
}

static void enableMoves(G_nodeList list) {
	for(G_nodeList p = list; p; p = p->tail) {
		G_node node = p->head;
		for(Live_moveList q = nodeMoves(node); q; q = q->tail) {
			Live_move move = q->head;
			Live_moveDList r = Move_look(moveHash, move);
			if(r->kind == ACT) {
				removeMoveFromSet(r);
				insertMoveToSet(workListMoves, r);
			}
		}
	}
}

static void decrementDegree(G_node node) {
	int nodeIndex = G_nodeIndex(node);
	if(degree[nodeIndex]-- == K) {
		enableMoves(G_NodeList(node, adjacent(node)));
		G_nodeDList p = G_look(nodeHash, node);
		removeNodeFromSet(p);
		if(moveRelated(node))
			insertNodeToSet(freezeWorkList, p);
		else
			insertNodeToSet(simplifyWorkList, p);
	}
}

static void addWorkList(G_node node) {
	G_nodeDList p = G_look(nodeHash, node);
	if(p->kind == PRECOL)
		return;
	int nodeIndex = G_nodeIndex(node);
	if(!moveRelated(node) && degree[nodeIndex] < K) {
		removeNodeFromSet(p);
		insertNodeToSet(simplifyWorkList, p);
	}
}

static bool ok(G_node t, G_node r) {
	// precondition: 'r' must be precolored
	int tIndex = G_nodeIndex(t);
	int rIndex = G_nodeIndex(r);
	G_nodeDList p = G_look(nodeHash, t);
	return degree[tIndex] < K || p->kind == PRECOL || adjSet[tIndex][rIndex];
}

static bool BriggsConservative(G_nodeList list) {
	int k = 0;
	for(G_nodeList p = list; p; p = p->tail) {
		G_node node = p->head;
		int nodeIndex = G_nodeIndex(node);
		k += (degree[nodeIndex] >= K); 
	}
	return k < K;
}

static bool GeorgeConservative(G_node u, G_node v) {
	// precondition: 'u' must be precolored
	for(G_nodeList p = adjacent(v); p; p = p->tail) {
		G_node t = p->head;
		if(!ok(t, u))
			return FALSE;
	}
	return TRUE;
}

static int getAlias(int nodeIndex) {
	G_node node = nodes[nodeIndex];
	G_nodeDList p = G_look(nodeHash, node);
	if(p->kind == COALNODE)
		return getAlias(alias[nodeIndex]);
	else
		return nodeIndex;
}

static void combine(G_node u, G_node v) {
	// precondition: 'u' may be precolored or live in either spillWorkList or
	// freezeWorkList while 'v' must live in either spillWorkList or 
	// freezeWorkList
	G_nodeDList up = G_look(nodeHash, u);
	G_nodeDList vp = G_look(nodeHash, v);
	removeNodeFromSet(vp);
	insertNodeToSet(coalescedNodes, vp);
	int uIndex = G_nodeIndex(u);
	int vIndex = G_nodeIndex(v);
	alias[vIndex] = uIndex;
	for(Live_moveList p = moveList[vIndex]; p; p = p->tail)
		insertMoveToList(uIndex, p->head);
	for(G_nodeList p = adjacent(v); p; p = p->tail) {
		G_node t = p->head;
		addEdge(t, u);
		decrementDegree(t);
	}
	if(degree[uIndex] >= K && up->kind == FREEWORK) {
		removeNodeFromSet(up);
		insertNodeToSet(spillWorkList, up);
	}
}

static void freezeMoves(G_node u) {
	// precondition: 'u' must be either low-degree and move-related or 
	// high-degree node
	int uIndex = G_nodeIndex(u);
	for(Live_moveList p = nodeMoves(u); p; p = p->tail) {
		Live_move move = p->head;
		int xIndex = getAlias(G_nodeIndex(move->src));
		int yIndex = getAlias(G_nodeIndex(move->dst));
		int vIndex = yIndex;
		if(uIndex == yIndex)
			vIndex = xIndex;
		G_node v = nodes[vIndex];
		// v' may be precolored or live in either spillWorkList or freezeWorkList
		// before freezing the move
		Live_moveDList mp = Move_look(moveHash, move);
		removeMoveFromSet(mp);
		insertMoveToSet(frozenMoves, mp);
		if(degree[vIndex] < K && !moveRelated(v)) {
			G_nodeDList vp = G_look(nodeHash, v);
			removeNodeFromSet(vp);
			insertNodeToSet(simplifyWorkList, vp);
		}
	}
}

static void simplify() {
	G_nodeDList p = simplifyWorkList->next;
	removeNodeFromSet(p);
	insertNodeToSet(selectStack, p);
	for(G_nodeList q = adjacent(p->node); q; q = q->tail)
		decrementDegree(q->head);
}

static void coalesce() {
	Live_moveDList p = workListMoves->next;
	Live_move move = p->move;
	int xIndex = getAlias(G_nodeIndex(move->src));
	int yIndex = getAlias(G_nodeIndex(move->dst));
	G_nodeDList yp = G_look(nodeHash, nodes[yIndex]);
	int uIndex = xIndex;
	int vIndex = yIndex;
	if(yp->kind == PRECOL) {
		uIndex = yIndex;
		vIndex = xIndex;
	}
	G_node u = nodes[uIndex];
	G_node v = nodes[vIndex];
	G_nodeDList up = G_look(nodeHash, u);
	G_nodeDList vp = G_look(nodeHash, v);
	removeMoveFromSet(p);
	if(u == v) {
		insertMoveToSet(coalescedMoves, p);
		addWorkList(u);
	}
	else if(vp->kind == PRECOL || adjSet[uIndex][vIndex]) {
		insertMoveToSet(constrainedMoves, p);
		addWorkList(u);
		addWorkList(v);
	}
	else if(up->kind == PRECOL && GeorgeConservative(u, v) ||
	        up->kind != PRECOL && BriggsConservative(mergeNodeList(adjacent(u), adjacent(v)))) {
		insertMoveToSet(coalescedMoves, p);
		combine(u, v);
		addWorkList(u);
	}
	else
		insertMoveToSet(activeMoves, p);
}

static void freeze() {
	G_nodeDList p = freezeWorkList->next;
	removeNodeFromSet(p);
	insertNodeToSet(simplifyWorkList, p);
	freezeMoves(p->node);
}

static void selectSpill() {
	G_nodeDList p = spillWorkList->next;
	removeNodeFromSet(p);
	insertNodeToSet(simplifyWorkList, p);
	freezeMoves(p->node);
}

static Temp_temp selectColor(Temp_tempList pigments, Temp_tempList exclude) {
	TAB_table tab = TAB_empty();
	for(Temp_tempList p = exclude; p; p = p->tail) 
		TAB_enter(tab, p->head, (void*)1);
	for(Temp_tempList p = pigments; p; p = p->tail) {
		if(!TAB_look(tab, p->head))
			return p->head;
	}
	return NULL;
}

static void assignColors(Temp_tempList pigments) {
	G_nodeDList np = selectStack->next, q;
	while(np->next) {
		q = np->next;
		removeNodeFromSet(np);
		G_node n = np->node;
		int nIndex = G_nodeIndex(n);
		Temp_tempList exclude = Temp_TempList(NULL, NULL);
		for(G_nodeList r = adjList[nIndex]; r; r = r->tail) {
			int wIndex = getAlias(G_nodeIndex(r->head));
			G_node w = nodes[wIndex];
			G_nodeDList wp = G_look(nodeHash, w);
			if(wp->kind == COLNODE || wp->kind == PRECOL)
				exclude->tail = Temp_TempList(color[wIndex], exclude->tail);
		}
		Temp_temp pigment = selectColor(pigments, exclude->tail);
		if(!pigment)  {
			insertNodeToSet(spilledNodes, np);
		}
		else {
			insertNodeToSet(coloredNodes, np);
			color[nIndex] = pigment;
			printf("color %s with %s\n", Temp_look(debug, Live_gtemp(n)), Temp_look(debug, color[nIndex]));
		}
		np = q;
	}
	for(np = coalescedNodes->next; np->next; np = np->next) {
		G_node n = np->node;
		int nIndex = G_nodeIndex(n);
		color[nIndex] = color[getAlias(nIndex)];
		printf("color %s with %s\n", Temp_look(debug, Live_gtemp(n)), Temp_look(debug, color[nIndex]));
	}	
}

static void init() {
	// initialize nodeSet
	preColored = G_NodeDList(NULL, PRECOL, NULL,
	              G_NodeDList(NULL, PRECOL, NULL, NULL));
	preColored->next->prev = preColored;
	initial = G_NodeDList(NULL, INITIAL, NULL,
	              G_NodeDList(NULL, INITIAL, NULL, NULL));
	initial->next->prev = initial;
	simplifyWorkList = G_NodeDList(NULL, SIMPWORK, NULL,
	              G_NodeDList(NULL, SIMPWORK, NULL, NULL));
	simplifyWorkList->next->prev = simplifyWorkList;
	freezeWorkList = G_NodeDList(NULL, FREEWORK, NULL,
	              G_NodeDList(NULL, FREEWORK, NULL, NULL));
	freezeWorkList->next->prev = freezeWorkList;
	spillWorkList = G_NodeDList(NULL, SPILWORK, NULL,
	              G_NodeDList(NULL, SPILWORK, NULL, NULL));
	spillWorkList->next->prev = spillWorkList;
	spilledNodes = G_NodeDList(NULL, SPILNODE, NULL,
	              G_NodeDList(NULL, SPILNODE, NULL, NULL));
	spilledNodes->next->prev = spilledNodes;
	coalescedNodes = G_NodeDList(NULL, COALNODE, NULL,
	              G_NodeDList(NULL, COALNODE, NULL, NULL));
	coalescedNodes->next->prev = coalescedNodes;
	coloredNodes = G_NodeDList(NULL, COLNODE, NULL,
	              G_NodeDList(NULL, COLNODE, NULL, NULL));
	coloredNodes->next->prev = coloredNodes;
	selectStack = G_NodeDList(NULL, STACK, NULL,
	              G_NodeDList(NULL, STACK, NULL, NULL));
	selectStack->next->prev = selectStack;
	
	// initialize moveSet
	coalescedMoves = Live_MoveDList(NULL, COAL, NULL,
	                  Live_MoveDList(NULL, COAL, NULL, NULL));
	coalescedMoves->next->prev = coalescedMoves;
	constrainedMoves = Live_MoveDList(NULL, CONS, NULL,
	                  Live_MoveDList(NULL, CONS, NULL, NULL));
	constrainedMoves->next->prev = constrainedMoves;
	frozenMoves = Live_MoveDList(NULL, FROZ, NULL,
	                  Live_MoveDList(NULL, FROZ, NULL, NULL));
	frozenMoves->next->prev = frozenMoves;
	workListMoves = Live_MoveDList(NULL, WORK, NULL,
	                  Live_MoveDList(NULL, WORK, NULL, NULL));
	workListMoves->next->prev = workListMoves;
	activeMoves = Live_MoveDList(NULL, ACT, NULL,
	                  Live_MoveDList(NULL, ACT, NULL, NULL));
	activeMoves->next->prev = activeMoves;

	// initialize array
	memset(adjSet, 0, sizeof(adjSet));	
	memset(adjList, 0, sizeof(adjList));
	memset(degree, 0, sizeof(degree));
	memset(moveList, 0, sizeof(moveList));
	memset(alias, 0, sizeof(alias));
	memset(color, 0, sizeof(color));
	memset(nodes, 0, sizeof(nodes));

	// initialize hash table
	nodeHash = G_empty();
	moveHash = Move_empty();	

	// initialize the number of machine register
	K = 0;
}

static void prepareNodes(G_graph ig, Temp_tempList regs) {
	for(G_nodeList p = G_nodes(ig); p; p = p->tail) {
		G_node node = p->head;
		int nodeIndex = G_nodeIndex(node);
		nodes[nodeIndex] = node;
		Temp_temp temp = Live_gtemp(node);
		bool hasColor = FALSE;
		for(Temp_tempList reg = regs; reg && !hasColor; reg = reg->tail) 
			hasColor = (temp == reg->head);
		G_nodeDList q = G_NodeDList(node, PRECOL, NULL, NULL);
		G_enter(nodeHash, node, q); // create mappings
		if(hasColor) {
			insertNodeToSet(preColored, q);
			color[nodeIndex] = temp;
			degree[nodeIndex] = 65536; // infinite
		}
		else
			insertNodeToSet(initial, q);
	} 
}

static void prepareAdjs(G_graph ig) {
	for(G_nodeList p = G_nodes(ig); p; p = p->tail) {
		G_node node = p->head;
		for(G_nodeList pred = G_pred(node); pred; pred = pred->tail) 
			addEdge(pred->head, node);
		for(G_nodeList succ = G_succ(node); succ; succ = succ->tail)
			addEdge(node, succ->head);
	}
}

static void prepareMoves(Live_moveList moves) {
	for(Live_moveList p = moves; p; p = p->tail) {
		Live_move move = p->head;
		G_node src = move->src;
		G_node dst = move->dst;
		int srcIndex = G_nodeIndex(src);
		int dstIndex = G_nodeIndex(dst);
		insertMoveToList(srcIndex, move);
		insertMoveToList(dstIndex, move);
		Live_moveDList q = Live_MoveDList(move, COAL, NULL, NULL);
		Move_enter(moveHash, move, q);
		insertMoveToSet(workListMoves, q);
	}
}

struct COL_result COL_color(struct Live_graph inter, Temp_map initial, Temp_tempList regs) {
	debug = Temp_layerMap(initial, Temp_name());
	init(); // initialize all the global variable
	// prepare preColored and initial node set and nodes[] and nodeHash
	prepareNodes(inter.graph, regs);
	prepareAdjs(inter.graph); // prepare adjSet[][], adjList[] and degree[]
	// prepare workListMoves move set and moveList[] and moveHash
	prepareMoves(inter.moves);
	// calculate the number of pigments
	for(Temp_tempList reg = regs; reg; reg = reg->tail, ++K);
	makeWorkList(); 
	while(TRUE) {
		if(!nodeSetEmpty(simplifyWorkList))
			simplify();
		else if(!moveSetEmpty(workListMoves))
			coalesce();
		else if(!nodeSetEmpty(freezeWorkList))
			freeze();
		else if(!nodeSetEmpty(spillWorkList))
			selectSpill();
		else
			break;
	}
	assignColors(regs);
	Temp_map coloring = Temp_empty();	
	for(G_nodeDList p = coloredNodes->next; p->next; p = p->next) {
		G_node node = p->node;
		int nodeIndex = G_nodeIndex(node);
		Temp_temp temp = Live_gtemp(node);
		Temp_enter(coloring, temp, Temp_look(initial, color[nodeIndex]));
	}
	for(G_nodeDList p = coalescedNodes->next; p->next; p = p->next) {
		G_node node = p->node;
		int nodeIndex = G_nodeIndex(node);
		Temp_temp temp = Live_gtemp(node);
		Temp_enter(coloring, temp, Temp_look(initial, color[nodeIndex]));
	}
	Temp_tempList spills = Temp_TempList(NULL, NULL);
	for(G_nodeDList p = spilledNodes->next; p->next; p = p->next) {
		G_node node = p->node;
		Temp_temp temp = Live_gtemp(node);
		spills->tail = Temp_TempList(temp, spills->tail);
	}
	struct COL_result result = {Temp_layerMap(initial, coloring), spills->tail};
	return result;
}

